# audio_manager.py
import io
import wave
import torchaudio
import librosa
import torch
import numpy as np
from pathlib import Path
import opuslib

# 音频处理配置 - 严格匹配官方参数
PROMPT_SAMPLE_RATE = 16000  # 提示音频必须16kHz
OUTPUT_SAMPLE_RATE = 24000  # 合成输出24kHz
TOP_DB = 60
HOP_LENGTH = 220
WIN_LENGTH = 440
MAX_VAL = 0.8  # 归一化最大值

# Opus编码配置
SAMPLE_RATE = 24000  # 采样率
CHANNELS = 1         # 声道数 (单声道)
FRAME_SIZE_MS = 20   # 每帧时长(毫秒)
FRAME_SIZE = int(SAMPLE_RATE * FRAME_SIZE_MS / 1000)  # 每帧样本数

# 音色类
class Voice:
    def __init__(self, name: str):
        self.name = name
        self.title = name
        self.prompt_text = f"我是{name}，这是我的声音。"
        self.audio_tensor = None
        self.file = None
        self.is_loaded = False

    def load_data(self, audio_file: Path):
        """使用官方处理流程加载提示音频"""
        try:
            # 加载并处理音频
            self.audio_tensor = load_and_process_prompt(str(audio_file))
            self.file = str(audio_file)
            
            # 尝试读取对应的文本文件
            text_file = audio_file.with_suffix(".txt")
            if text_file.exists():
                self._parse_text_file(text_file)
            
            self.is_loaded = True
            print(f"✅ 加载音色: {self.name} ({self.audio_tensor.shape})")
            return True
        except Exception as e:
            print(f"❌ 加载音色失败 {self.name}: {str(e)}")
            return False

    def _parse_text_file(self, text_path: Path):
        """修复文本解析逻辑"""
        try:
            content = text_path.read_text(encoding='utf-8').strip()
        except:
            try:
                content = text_path.read_text(encoding='gbk').strip()
            except:
                return
        
        if not content:
            return
        
        # 尝试从文本中提取标题和提示语
        lines = content.split('\n')
        if len(lines) > 0:
            self.title = lines[0].strip()
        if len(lines) > 1:
            self.prompt_text = lines[1].strip()
        else:
            self.prompt_text = self.title
        
        print(f"📝 加载描述: {self.name} -> {self.prompt_text}")

# 加载提示音频
def load_and_process_prompt(file_path):
    """
    严格遵循官方流程处理提示音频：
    1. 加载音频
    2. 转换为单声道
    3. 重采样到16kHz
    4. 应用标准化处理
    """
    # 使用torchaudio加载
    speech, sr = torchaudio.load(file_path)
    
    # 转为单声道
    if speech.shape[0] > 1:
        speech = speech.mean(0, keepdim=True)
    
    # 重采样到模型需要的16kHz
    if sr != PROMPT_SAMPLE_RATE:
        resampler = torchaudio.transforms.Resample(
            orig_freq=sr, new_freq=PROMPT_SAMPLE_RATE
        )
        speech = resampler(speech)
    
    # 应用官方处理流程
    return official_audio_processing(speech, sample_rate=PROMPT_SAMPLE_RATE)

# 官方精确复制版音频处理功能
def official_audio_processing(audio_data, sample_rate=OUTPUT_SAMPLE_RATE):
    """
    精确复制官方音频处理流程：
    1. 格式处理
    2. 静音修剪
    3. 音量归一化
    """
    # 1. 转换为PyTorch张量
    if isinstance(audio_data, np.ndarray):
        speech = torch.from_numpy(audio_data).float()
    else:
        speech = audio_data.float()
    
    # 2. 单声道处理
    if len(speech.shape) > 1:
        if speech.shape[0] > 1:
            speech = speech.mean(0, keepdim=True)
        speech = speech.squeeze(0)
    
    # 3. 静音修剪
    if speech.numel() > 0:  # 确保非空
        speech = speech.unsqueeze(0)
        try:
            speech_np = speech.squeeze(0).numpy()
            trimmed_speech, _ = librosa.effects.trim(
                speech_np,
                top_db=TOP_DB,
                frame_length=WIN_LENGTH,
                hop_length=HOP_LENGTH
            )
            speech = torch.tensor(trimmed_speech).float().unsqueeze(0)
        except Exception as e:
            print(f"静音修剪失败: {str(e)}")
    
    # 4. 音量归一化
    max_amp = speech.abs().max()
    if max_amp > MAX_VAL and max_amp > 0:
        speech = speech / max_amp * MAX_VAL
    
    return speech
    
## PCM处理工具
def float32_to_int16(float_array):
    """将 float32 [-1, 1] 转换为 int16"""
    # 确保数据在有效范围内
    float_array = np.clip(float_array, -1.0, 1.0)
    # 缩放并转换为 int16
    return (float_array * 32767).astype(np.int16)

## Opus编码处理
def create_opus_encoder():
    """创建 Opus 编码器"""
    enc = opuslib.Encoder(SAMPLE_RATE, CHANNELS, opuslib.APPLICATION_AUDIO)
    # 设置关键编码参数
    enc.bitrate = 24000  # 比特率 (bps)
    enc.complexity = 6   # 复杂度 (1-10)
    enc.signal = opuslib.SIGNAL_VOICE
    return enc

def float32_pcm_to_opus(float32_pcm: np.ndarray, encoder=None) -> bytes:
    """
    将float32 PCM数据编码为Opus
    :param float32_pcm: float32格式的PCM数据 ([-1, 1]范围)
    :param encoder: 已创建的Opus编码器 (可选)
    :return: Opus编码的音频数据
    """
    # 转换为int16
    int16_data = float32_to_int16(float32_pcm)
    return int16_pcm_to_opus(int16_data, encoder)

def int16_pcm_to_opus(int16_pcm: np.ndarray, encoder=None) -> bytes:
    """
    将int16 PCM数据编码为Opus
    :param int16_pcm: int16格式的PCM数据
    :param encoder: 已创建的Opus编码器 (可选)
    :return: Opus编码的音频数据
    """
    # 创建编码器（如果需要）
    enc = encoder or create_opus_encoder()
    
    # 将数据平坦化为1D数组
    if len(int16_pcm.shape) > 1:
        int16_pcm = int16_pcm.flatten()
        
    # 编码为Opus
    return enc.encode(int16_pcm.tobytes(), FRAME_SIZE)

## WAV编码处理
def float32_pcm_to_wav(float32_pcm: np.ndarray, sample_rate: int = 24000) -> bytes:
    """
    将 float32 PCM 数据转换为 WAV 格式的内存流
    :param float32_pcm: float32格式的PCM数据 ([-1, 1]范围)
    :param sample_rate: 采样率 (默认 24000 Hz)
    :return: WAV 音频字节流
    """
    # 转换为int16
    int16_data = float32_to_int16(float32_pcm)
    return int16_pcm_to_wav(int16_data, sample_rate)

def int16_pcm_to_wav(int16_data: np.ndarray, sample_rate: int = 24000) -> bytes:
    """
    将int16 PCM数据转换为WAV格式
    :param int16_data: int16格式的PCM数据
    :param sample_rate: 采样率 (默认24000 Hz)
    :return: WAV音频字节流
    """
    # 确保是单声道
    if len(int16_data.shape) > 1 and int16_data.shape[0] != 1:
        int16_data = int16_data.flatten()
    
    # 创建内存中的二进制流
    with io.BytesIO() as wav_buffer:
        with wave.open(wav_buffer, 'wb') as wav_file:
            wav_file.setnchannels(1)  # 单声道
            wav_file.setsampwidth(2)  # 16位 = 2字节
            wav_file.setframerate(sample_rate)
            wav_file.setcomptype('NONE', 'not compressed')
            wav_file.writeframes(int16_data.tobytes())
            
        # 返回完整的WAV字节流
        return wav_buffer.getvalue()